{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1.import libs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"A script for data collection on Quafu platform.\"\"\"\n",
    "import sys\n",
    "import os\n",
    "sys.path.append(os.path.abspath(os.path.join(os.getcwd(), '..')))\n",
    "\n",
    "from utils.get_data import ShotTestReceiver, CircuitReceiver\n",
    "from qiskit import QuantumCircuit, ClassicalRegister\n",
    "from qiskit.qasm2 import dumps, loads  # Use this for QASM conversion\n",
    "import itertools\n",
    "import random"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2.define some helpful functions\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_readout_benchmark_circuit(qubit_indices: list, use_QED: bool=False, use_QED_reverse: bool=False, QED_map: dict=None, inherit_circuit: str=None) -> list[str]:\n",
    "    \"\"\"Create benchmark circuits for readout error characterization.\n",
    "    \n",
    "    This function generates circuits that prepare all possible computational basis\n",
    "    states for the specified qubits to characterize readout errors.\n",
    "    \n",
    "    Args:\n",
    "        qubit_indices: List of physical qubit indices to benchmark\n",
    "        use_QED: bool, whether to use QED to encode the state\n",
    "        use_QED_reverse: bool, whether to use reverse QED to encode the state\n",
    "        QED_map: dict, the mapping of the qubits to the QED qubits, e.g. {37: [44]}, means the qubit 37 is mapped to the qubit 44 in the QED\n",
    "        \n",
    "    Returns:\n",
    "        list[str]: List of QASM strings for all state preparation combinations\n",
    "    \"\"\"\n",
    "    n_qubits = len(qubit_indices)\n",
    "    benchmark_n = len(qubit_indices)\n",
    "    if QED_map is not None:\n",
    "        benchmark_n = len(QED_map)\n",
    "    preparations = list(itertools.product([0,1], repeat=benchmark_n))\n",
    "    circuits = []\n",
    "    \n",
    "    for prep in preparations:\n",
    "        max_qubit_idx = max(qubit_indices) + 1\n",
    "        if inherit_circuit is not None:\n",
    "            circuit = loads(inherit_circuit)\n",
    "        else:\n",
    "            circuit = QuantumCircuit(max_qubit_idx, n_qubits)\n",
    "        \n",
    "        # Prepare basis states\n",
    "        for qubit, state in zip(qubit_indices, prep):\n",
    "            if state == 1:\n",
    "                circuit.x(qubit)\n",
    "\n",
    "        if use_QED:\n",
    "            for qubit, map_qubits in QED_map.items():\n",
    "                for map_qubit in map_qubits:\n",
    "                    circuit.cx(qubit, map_qubit)\n",
    "        if use_QED_reverse:\n",
    "            for qubit, map_qubits in QED_map.items():\n",
    "                for map_qubit in map_qubits:\n",
    "                    circuit.cx(qubit, map_qubit)\n",
    "                    circuit.x(map_qubit)\n",
    "\n",
    "        # Add measurements\n",
    "        for i, qubit in enumerate(qubit_indices):\n",
    "            circuit.measure(qubit, i)\n",
    "            \n",
    "        circuits.append(dumps(circuit))\n",
    "        \n",
    "    return circuits"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3.get data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Execution date: 2025-04-09\n"
     ]
    }
   ],
   "source": [
    "# Get the current date when this code is executed\n",
    "from datetime import datetime\n",
    "date = datetime.now().strftime(\"%Y-%m-%d\")\n",
    "print(f\"Execution date: {date}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3.1 get benchmark data(2-qubit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "qubit_indices = [44, 50]    # readout error rate 1.38 1.14\n",
    "# qubit_indices = [20, 26]    # readout error rate 1.12 1.26\n",
    "# qubit_indices = [2, 9]    # readout error rate 1.67 1.6\n",
    "# qubit_indices = [23, 29]    # readout error rate 2.23 2.39"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "benchmark_type = 'benchmark'\n",
    "n_qubits = len(qubit_indices)\n",
    "shot_test_platform_config = {\n",
    "    'platform': 'tianyan',\n",
    "    # 'login_key': '8N8TOqEEhrGznQPEoguaiDF6NEFWjSdL2UlZwDV9Ks8.Qf3YjMyQTNyQzNxojIwhXZiwiM4kDN6ICZpJye.9JiN1IzUIJiOicGbhJCLiQ1VKJiOiAXe0Jye',\n",
    "    'login_key': 'kCO7NfSBK79mgrEV+zE39oeGLBVrY+JUBQ07lp22NL0=',\n",
    "    'machine_name': 'tianyan176-2',\n",
    "    'num_qubits': n_qubits,\n",
    "    'exp_name': f'{len(qubit_indices)}-qubit_qubits_{\"-\".join(map(str, qubit_indices))}'  # Include number of qubits, circuit type and qubit indices\n",
    "}\n",
    "\n",
    "# shot_range = [int(1e3), 5*int(1e3), int(1e4)] #, int(1e5)]    # different shots\n",
    "shot_range = [int(1e4)]     # same shots\n",
    "benchmark_receiver = ShotTestReceiver(shot_test_platform_config, shot_range=shot_range, circuit_nums=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Job submission completed.\n"
     ]
    }
   ],
   "source": [
    "benchmark_receiver.submit_experiment(create_readout_benchmark_circuit(qubit_indices=qubit_indices))\n",
    "\n",
    "benchmark_receiver.receive_results()\n",
    "benchmark_receiver.save_results(f'../results/{date}/{benchmark_type}/')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3.2 get benchmark data(1-qubit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "# qubit_indices = [44]    # readout error rate 1.38 1.14\n",
    "# qubit_indices = [50]    # readout error rate 1.38 1.14\n",
    "# qubit_indices = [20]    # readout error rate 1.12 1.26\n",
    "# qubit_indices = [26]    # readout error rate 1.12 1.26\n",
    "# qubit_indices = [2]    # readout error rate 1.67 1.6\n",
    "qubit_indices = [9]    # readout error rate 1.67 1.6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "benchmark_type = 'benchmark'\n",
    "n_qubits = len(qubit_indices)\n",
    "shot_test_platform_config = {\n",
    "    'platform': 'tianyan',\n",
    "    # 'login_key': '8N8TOqEEhrGznQPEoguaiDF6NEFWjSdL2UlZwDV9Ks8.Qf3YjMyQTNyQzNxojIwhXZiwiM4kDN6ICZpJye.9JiN1IzUIJiOicGbhJCLiQ1VKJiOiAXe0Jye',\n",
    "    'login_key': 'kCO7NfSBK79mgrEV+zE39oeGLBVrY+JUBQ07lp22NL0=',\n",
    "    'machine_name': 'tianyan176-2',\n",
    "    'num_qubits': n_qubits,\n",
    "    'exp_name': f'{len(qubit_indices)}-qubit_qubits_{\"-\".join(map(str, qubit_indices))}'  # Include number of qubits, circuit type and qubit indices\n",
    "}\n",
    "\n",
    "# shot_range = [int(1e3), 5*int(1e3), int(1e4)] #, int(1e5)]    # different shots\n",
    "shot_range = [int(1e4)]     # same shots\n",
    "benchmark_receiver = ShotTestReceiver(shot_test_platform_config, shot_range=shot_range, circuit_nums=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Job submission completed.\n"
     ]
    }
   ],
   "source": [
    "benchmark_receiver.submit_experiment(create_readout_benchmark_circuit(qubit_indices=qubit_indices))\n",
    "\n",
    "benchmark_receiver.receive_results()\n",
    "benchmark_receiver.save_results(f'../results/{date}/{benchmark_type}/')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
